import numpy as np
import plotly.graph_objects as go
from plotly.offline import plot

# 原子半径
r0 = 0.55;
r1 = 0.45;

# 连接中心原子的原子或电子对个数个数. N_PAIRS=0,1,2,3,4
N_PAIRS=4;
# 其中孤对电子个数. N_ELECTRONS=0,1,...,N_PAIRS
N_ELECTRONS=1;

def add_sphere(x,y,z,cx,cy,cz,r):
    cu = np.sqrt((x-cx)**2+(y-cy)**2+(z-cz)**2)*(2/r)
    return np.exp(-cu)*(2/r)

Atom = np.zeros((5, 3))
isTransparent = np.zeros((5, 1))
a = 0

if N_PAIRS == 0:
    Atom[0, :] = [0, 0, 0]
elif N_PAIRS == 1:
    a = r0 + r1
    Atom[0, :] = [0, 0, 0]
    Atom[1, :] = [a, 0, 0]
elif N_PAIRS == 2:
    a = r0 + r1
    Atom[0, :] = [0, 0, 0]
    Atom[1, :] = [a, 0, 0]
    Atom[2, :] = [-a, 0, 0]
elif N_PAIRS == 3:
    a = np.sqrt(3) * (r0 + r1)
    Atom[0, :] = [a / 2, np.sqrt(3) / 6 * a, 0]
    Atom[1, :] = [0, 0, 0]
    Atom[2, :] = [a, 0, 0]
    Atom[3, :] = [a / 2, np.sqrt(3) / 2 * a, 0]
elif N_PAIRS == 4:
    a = np.sqrt(8 / 3) * (r0 + r1)
    Atom[0, :] = [a / 2, np.sqrt(3) / 6 * a, np.sqrt(6) / 12 * a]
    Atom[1, :] = [0, 0, 0]
    Atom[2, :] = [a, 0, 0]
    Atom[3, :] = [a / 2, np.sqrt(3) / 2 * a, 0]
    Atom[4, :] = [a / 2, np.sqrt(3) / 6 * a, np.sqrt(6) / 3 * a]
else:
    raise ValueError('invalid N_PAIRS')

# 平移到原点
for i in range(1,N_PAIRS+1):
  Atom[i,:] -= Atom[0,:];
Atom[0,:] = [0,0,0];

if N_ELECTRONS > N_PAIRS:
    raise ValueError('invalid N_ELECTRONS.')

if N_ELECTRONS == 0:
    pass
else:
    isTransparent[N_PAIRS-N_ELECTRONS+1:N_PAIRS+1] = 1;

# Create Mesh
L = 2;
dx = 0.1
x,y,z=np.meshgrid(np.arange(-L,L,dx),np.arange(-L,L,dx),np.arange(-L,L,dx))
x,y,z=x.flatten(),y.flatten(),z.flatten()

# Voxel Data
u = np.zeros_like(x)

for k in range(N_PAIRS+1):
    r = r0 if i == 0 else r1;
    c = 0.5 if isTransparent[k] == 1 else 1;
    u += c*add_sphere(x,y,z,Atom[k,0],Atom[k,1],Atom[k,2],r)

fig = go.Figure(data=go.Volume(
    x=x,
    y=y,
    z=z,
    value=u,
    opacity=0.1, # needs to be small to see through all surfaces
    surface_count=20, # needs to be a large number for good volume rendering
    ))
plot(fig)
